package com.TripMate.system.AI;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

@Component
public class HttpSSE {
    private static final Logger log = LoggerFactory.getLogger(HttpSSE.class);
    
    // 阿里云DeepSeek V3模型API密钥
    private static final String API_KEY = "sk-5f5b9c53dfed4749a7021bb5c13274d5";
    // DeepSeek V3模型ID
    private static final String MODEL_ID = "deepseek-v3";

    public CompletableFuture<String> sseInvoke(String content, String sid, String requestId) {
        CompletableFuture<String> future = new CompletableFuture<>();
        AtomicReference<StringBuilder> responseBuilder = new AtomicReference<>(new StringBuilder());

        // 创建一个线程来处理流式调用，避免阻塞主线程
        ExecutorService executor = Executors.newSingleThreadExecutor();
        executor.submit(() -> {
            try {
                // 创建生成请求
                Generation gen = new Generation();
                
                // 创建用户消息
                Message userMsg = Message.builder()
                        .role(Role.USER.getValue())
                        .content(content)
                        .build();
                
                // 构建参数
                GenerationParam param = GenerationParam.builder()
                        .apiKey(API_KEY)
                        .model(MODEL_ID)
                        .messages(Arrays.asList(userMsg))
                        .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                        .incrementalOutput(true)
                        .build();
                
                // 流式调用，使用blockingForEach处理每个结果
                Flowable<GenerationResult> result = gen.streamCall(param);
                result.blockingForEach(message -> {
                    try {
                        // 获取内容
                        String content1 = message.getOutput().getChoices().get(0).getMessage().getContent();
                        if (content1 != null && !content1.isEmpty()) {
                            log.debug("Received content: {}", content1);
                            responseBuilder.get().append(content1);
                        }
                    } catch (Exception e) {
                        log.error("Error processing message: ", e);
                    }
                });
                
                // 流式调用完成后，返回完整响应
                String finalResponse = responseBuilder.get().toString().trim();
                future.complete(finalResponse);
                
            } catch (NoApiKeyException | ApiException | InputRequiredException e) {
                log.error("API error:", e);
                future.completeExceptionally(e);
            } catch (Exception e) {
                log.error("Unexpected error:", e);
                future.completeExceptionally(e);
            } finally {
                executor.shutdown();
            }
        });

        // 设置超时处理
        ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
        scheduler.schedule(() -> {
            if (!future.isDone()) {
                String response = responseBuilder.get().toString();
                if (!response.isEmpty()) {
                    future.complete(response);
                } else {
                    future.completeExceptionally(new TimeoutException("SSE request timed out"));
                }
                executor.shutdownNow();
            }
            scheduler.shutdown();
        }, 300, TimeUnit.SECONDS);

        return future;
    }
}
